package com.zwh.common.utils.servlet;

import java.io.*;
import java.net.URLEncoder;
import java.nio.charset.Charset;
import java.util.*;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.ServletRequest;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpMethod;
import org.springframework.lang.Nullable;



/**
 * 自定义request包装，避免request.getInputStream只能被读取一次就关闭的问题
 * 参考：org.springframework.web.util.ContentCachingRequestWrapper
 * @author zwh
 * @date 2020/12/9 15:22
 **/

public class MyCachingRequestWrapper extends HttpServletRequestWrapper {
    Logger log = LoggerFactory.getLogger(MyCachingRequestWrapper.class);
    private static final String FORM_CONTENT_TYPE = "application/x-www-form-urlencoded";
    private final ByteArrayOutputStream cachedContent;
    @Nullable
    private final Integer contentCacheLimit;

    /**
     * 缓存inputStream
     */
    @Nullable
    private ServletInputStream inputStream;
    /**
     * inputStream包装Reader类
     */
    @Nullable
    private BufferedReader reader;
    /**
     * 存储body数据的容器
     */
//    private final byte[] body;

    public MyCachingRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        int contentLength = request.getContentLength();
        this.cachedContent = new ByteArrayOutputStream(contentLength >= 0 ? contentLength : 1024);
        this.contentCacheLimit = null;
        // 将body数据存储起来
//        String bodyStr = getBodyString(request);
//        body = bodyStr.getBytes(Charset.defaultCharset());
    }

    /**
     * 获取请求Body
     *
     * @param request request
     * @return String
     */
    public String getBodyString(final ServletRequest request) {
        try {
            return inputStream2String(request.getInputStream());
        } catch (IOException e) {
            log.error("", e);
            throw new RuntimeException(e);
        }
    }

    /**
     * 获取请求Body
     *
     * @return String
     */
    public String getBodyString() throws IOException {
        return inputStream2String(getInputStream());
    }

    /**
     * 将inputStream里的数据读取出来并转换成字符串
     *
     * @param inputStream inputStream
     * @return String
     */
    public String inputStream2String(InputStream inputStream) {
        StringBuilder sb = new StringBuilder();
        BufferedReader reader = null;

        try {
            reader = new BufferedReader(new InputStreamReader(inputStream, Charset.defaultCharset()));
            String line;
            while ((line = reader.readLine()) != null) {
                sb.append(line);
            }
        } catch (IOException e) {
            log.error("", e);
            throw new RuntimeException(e);
        } finally {
            if (reader != null) {
                try {
                    reader.close();
                } catch (IOException e) {
                    log.error("", e);
                }
            }
        }

        return sb.toString();
    }


    @Override
    public ServletInputStream getInputStream() throws IOException {

        if (this.inputStream == null) {
            this.inputStream = new MyCachingRequestWrapper.ContentCachingInputStream(getRequest().getInputStream());
        }
        return this.inputStream;
    }
    @Override
    public BufferedReader getReader() throws IOException {
        if (this.reader == null) {
            this.reader = new BufferedReader(new InputStreamReader(getInputStream(), getCharacterEncoding()));
        }
        return this.reader;
    }
    @Override
    public String getCharacterEncoding() {
        String enc = super.getCharacterEncoding();
        return (enc != null ? enc : Charset.defaultCharset().name());
    }

    @Override
    public String getParameter(String name) {
        if (this.cachedContent.size() == 0 && isFormPost()) {
            writeRequestParametersToCachedContent();
        }
        return super.getParameter(name);
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        if (this.cachedContent.size() == 0 && isFormPost()) {
            writeRequestParametersToCachedContent();
        }
        return super.getParameterMap();
    }

    @Override
    public Enumeration<String> getParameterNames() {
        if (this.cachedContent.size() == 0 && isFormPost()) {
            writeRequestParametersToCachedContent();
        }
        return super.getParameterNames();
    }

    @Override
    public String[] getParameterValues(String name) {
        if (this.cachedContent.size() == 0 && isFormPost()) {
            writeRequestParametersToCachedContent();
        }
        return super.getParameterValues(name);
    }


    private boolean isFormPost() {
        String contentType = getContentType();
        return (contentType != null && contentType.contains(FORM_CONTENT_TYPE) &&
                HttpMethod.POST.matches(getMethod()));
    }

    private void writeRequestParametersToCachedContent() {
        try {
            if (this.cachedContent.size() == 0) {
                String requestEncoding = getCharacterEncoding();
                Map<String, String[]> form = super.getParameterMap();
                for (Iterator<String> nameIterator = form.keySet().iterator(); nameIterator.hasNext(); ) {
                    String name = nameIterator.next();
                    List<String> values = Arrays.asList(form.get(name));
                    for (Iterator<String> valueIterator = values.iterator(); valueIterator.hasNext(); ) {
                        String value = valueIterator.next();
                        this.cachedContent.write(URLEncoder.encode(name, requestEncoding).getBytes());
                        if (value != null) {
                            this.cachedContent.write('=');
                            this.cachedContent.write(URLEncoder.encode(value, requestEncoding).getBytes());
                            if (valueIterator.hasNext()) {
                                this.cachedContent.write('&');
                            }
                        }
                    }
                    if (nameIterator.hasNext()) {
                        this.cachedContent.write('&');
                    }
                }
            }
        }
        catch (IOException ex) {
            throw new IllegalStateException("Failed to write request parameters to cached content", ex);
        }
    }

    /**
     * Return the cached request content as a byte array.
     * <p>The returned array will never be larger than the content cache limit.
     * @see #ContentCachingRequestWrapper(HttpServletRequest, int)
     */
    public byte[] getContentAsByteArray() {
        return this.cachedContent.toByteArray();
    }
    /**
     * Template method for handling a content overflow: specifically, a request
     * body being read that exceeds the specified content cache limit.
     * <p>The default implementation is empty. Subclasses may override this to
     * throw a payload-too-large exception or the like.
     * @param contentCacheLimit the maximum number of bytes to cache per request
     * which has just been exceeded
     * @since 4.3.6
     * @see #ContentCachingRequestWrapper(HttpServletRequest, int)
     */
    protected void handleContentOverflow(int contentCacheLimit) {
    }


    /**
     * inputStream包装类，读取时会缓存内容
     */
    private class ContentCachingInputStream extends ServletInputStream {

        private final ServletInputStream is;

        private boolean overflow = false;

        public ContentCachingInputStream(ServletInputStream is) {
            this.is = is;
        }

        /**
         * 从输入流中读取数据的下一个字节，返回0-255范围内的int字节数，
         * 如果到达流末尾时，返回-1，表示数据已经读完
         * 这个方法处理效率较低，不常用
         */
        @Override
        public int read() throws IOException {
            int ch = this.is.read();
            if (ch != -1 && !this.overflow) {
                if (contentCacheLimit != null && cachedContent.size() == contentCacheLimit) {
                    this.overflow = true;
                    handleContentOverflow(contentCacheLimit);
                }
                else {
                    cachedContent.write(ch);
                }
            }
            return ch;
        }

        @Override
        public int read(byte[] b) throws IOException {
            int count = this.is.read(b);
            writeToCache(b, 0, count);
            return count;
        }

        /**
         * 将输入流中最多 len 个数据字节读入 byte 数组。<br/>
         * 尝试读取len 个字节，但读取的字节也可能小于该值。以整数形式返回实际读取的字节数
         * @param b
         * @param off
         * @param len
         * @return
         * @throws IOException
         */
        @Override
        public int read(final byte[] b, final int off, final int len) throws IOException {
            int count = this.is.read(b, off, len);
            writeToCache(b, off, count);
            return count;
        }

        @Override
        public int readLine(final byte[] b, final int off, final int len) throws IOException {
            int count = this.is.readLine(b, off, len);
            writeToCache(b, off, count);
            return count;
        }

        private void writeToCache(final byte[] b, final int off, int count) {
            if (!this.overflow && count > 0) {
                if (contentCacheLimit != null &&
                        count + cachedContent.size() > contentCacheLimit) {
                    this.overflow = true;
                    cachedContent.write(b, off, contentCacheLimit - cachedContent.size());
                    handleContentOverflow(contentCacheLimit);
                    return;
                }
                cachedContent.write(b, off, count);
            }
        }

        /**
         * 当前流中所有数据读取完成时返回true，否则返回false
         * @return
         */
        @Override
        public boolean isFinished() {
            return this.is.isFinished();
        }

        /**
         * 如果可以在不阻塞的情况下读取数据，则返回true，否则返回false
         * @return
         */
        @Override
        public boolean isReady() {
            return this.is.isReady();
        }

        /**
         * 为这个ServletInputStream设置ReadListener，从而切换到非阻塞IO。只有在异步处理或HTTP升级处理中切换到非阻塞IO才有效。
         * @param readListener
         */
        @Override
        public void setReadListener(ReadListener readListener) {
            this.is.setReadListener(readListener);
        }
    }
}
